# var_multivar_4day.py
import argparse, warnings
warnings.filterwarnings("ignore")
import pandas as pd
import numpy as np
from statsmodels.tsa.api import VAR
from pandas.tseries.offsets import MonthEnd

TRAIN_YEARS = list(range(2017, 2023))
FORE_YEARS  = [2025, 2026]

def make_4day_grid_for_year(year: int):
    d = pd.Timestamp(year=year, month=3, day=1)
    end = pd.Timestamp(year=year, month=7, day=1)
    out = []
    while d <= end:
        out.append(d); d += pd.Timedelta(days=4)
    return pd.DatetimeIndex(out)

def build_future_grid():
    return pd.DatetimeIndex(sorted(np.concatenate([make_4day_grid_for_year(y) for y in FORE_YEARS])))

def ensure_percent_scale(s: pd.Series) -> pd.Series:
    s = pd.to_numeric(s, errors="coerce")
    if s.max(skipna=True) <= 1.0: s = s * 100.0
    return s.clip(lower=0, upper=100)

def load_train_matrix(path, target_col, cols=None):
    df = pd.read_csv(path)
    # 识别时间
    date_col = None
    for c in ["date","时间","日期","ds","time"]:
        if c in df.columns: date_col = c; break
    if date_col is None:
        raise ValueError("训练表未找到时间列（date/时间/日期/ds/time）")
    df["date"] = pd.to_datetime(df[date_col])
    # 仅保留 2017–2022 的训练行
    df = df[(df["date"].dt.year>=min(TRAIN_YEARS)) & (df["date"].dt.year<=max(TRAIN_YEARS))].copy()
    # 目标列
    if target_col not in df.columns:
        raise ValueError(f"找不到目标列 {target_col}")
    df[target_col] = ensure_percent_scale(df[target_col])
    # 选择回归因子
    if cols is None:
        # 默认使用训练表中除时间/目标外的所有数值列
        num_cols = df.select_dtypes(include=[np.number]).columns.tolist()
        cols = [c for c in num_cols if c != target_col]
    # 组成内生矩阵：目标列放第一列便于输出
    use_cols = [target_col] + cols
    X = df[["date"] + use_cols].copy().set_index("date").sort_index()
    # 确保是 4 天等距栅格（如果预处理表就是该栅格，这步只是在防御式补齐）
    full_idx = pd.date_range(X.index.min(), X.index.max(), freq="4D")
    X = X.reindex(full_idx).interpolate("time").ffill().bfill()
    return X, use_cols

def var_forecast_level(X: pd.DataFrame, steps: int):
    """对一阶差分建模，预测差分后再积分回到水平"""
    # 一阶差分
    dX = X.diff().dropna()
    # 按 AIC 选 lag（最多 12 个 lag）
    sel = VAR(dX).select_order(12)
    p = 3
    try:
        if hasattr(sel, "selected_orders") and sel.selected_orders:
            p = sel.selected_orders.get("aic", p) or p
    except Exception:
        pass
    model = VAR(dX).fit(maxlags=p, ic=None)
    fc_diff = model.forecast(dX.values[-p:], steps=steps)
    fc_diff = pd.DataFrame(fc_diff, columns=dX.columns)
    # 累加还原水平
    last_level = X.iloc[-1]
    fc_levels = fc_diff.cumsum().add(last_level.values, axis=1)
    # 构造时间索引（4天步长）
    future_idx = pd.date_range(X.index[-1] + pd.Timedelta(days=4), periods=steps, freq="4D")
    fc_levels.index = future_idx
    return fc_levels

def main():
    ap = argparse.ArgumentParser(description="VAR 多元预测（3/1~7/1 每4天）")
    ap.add_argument("--train", required=True, help="训练表 CSV（含就业率+回归因子，建议用 pre_2016_2022_grid.csv）")
    ap.add_argument("--target-col", default="employment_rate")
    ap.add_argument("--cols", nargs="*", default=None, help="参与建模的因子列（至少4个）")
    ap.add_argument("--out", default="var_forecast_2025_2026_4day.csv")
    args = ap.parse_args()

    X, col_order = load_train_matrix(args.train, args.target_col, args.cols)

    # 需要预测的未来步数（两年的 4天栅格长度）
    fut_idx = build_future_grid()
    steps = len(fut_idx)

    fc_levels = var_forecast_level(X, steps=steps)

    # 取目标列（第一列）并对齐到未来栅格（两者长度一致，直接重设索引）
    yhat = fc_levels.iloc[:, 0].copy()
    yhat = pd.Series(yhat.values, index=fut_idx)
    out = pd.DataFrame({"date": yhat.index, "就业率（%）": yhat.clip(0,100).round(6)})
    out.to_csv(args.out, index=False, encoding="utf-8-sig")
    print(f"[VAR] saved: {args.out} rows={len(out)}")

if __name__ == "__main__":
    main()
